package net.lucene.chapter5.analyzer;

import com.chenlb.mmseg4j.ComplexSeg;
import com.chenlb.mmseg4j.Dictionary;
import com.chenlb.mmseg4j.Seg;
import com.chenlb.mmseg4j.analysis.MMSegTokenizer;
import org.apache.lucene.analysis.Analyzer;
import org.apache.lucene.analysis.TokenFilter;
import org.apache.lucene.analysis.TokenStream;
import org.apache.lucene.analysis.tokenattributes.CharTermAttribute;
import org.apache.lucene.analysis.tokenattributes.PositionIncrementAttribute;
import org.apache.lucene.util.AttributeSource;

import java.io.File;
import java.io.IOException;
import java.io.Reader;
import java.util.HashMap;
import java.util.Map;
import java.util.Stack;

/**
 * 自定义的同义词分词器
 *
 * @author Barudisshu
 */
public final class MySynonymAnalyzer extends Analyzer {

    // 存放词汇文件目录
    private static final String PATH_OF_DATA = "chapter5-demo/src/main/resources/lucene/data";

    @Override
    public TokenStream tokenStream(String fieldName, Reader reader) {
        return new MySynonymTokenFilter(new MyMMSegTokenizer(new ComplexSeg(Dictionary.getInstance(new File(PATH_OF_DATA))), reader));
    }

    /**
     * 自定义的TokenFilter
     */
    final class MySynonymTokenFilter extends TokenFilter {
        private CharTermAttribute cta;              // 用于获取TokenStream中的语汇单元
        private PositionIncrementAttribute pia;     // 用于获取TokenStream中的位置增量
        private AttributeSource.State tokenState;   // 用于保存语汇单元的状态
        private Stack<String> synonymStack;         // 用于保存同义词

        protected MySynonymTokenFilter(TokenStream input) {
            super(input);
            this.cta = this.addAttribute(CharTermAttribute.class);
            this.pia = this.addAttribute(PositionIncrementAttribute.class);
            this.synonymStack = new Stack<>();
        }

        /**
         * 判断是否存在同义词
         */
        private boolean isHaveSynonym(String name) {
            // 先定义同义词的词典
            Map<String, String[]> synonymMap = new HashMap<>();
            synonymMap.put("我", new String[]{"咱", "俺"});
            synonymMap.put("中国", new String[]{"兲朝", "大陆"});
            if (synonymMap.containsKey(name)) {
                for (String str : synonymMap.get(name)) {
                    this.synonymStack.push(str);
                }
                return true;
            }
            return false;
        }

        @Override
        public boolean incrementToken() throws IOException {

            if (this.synonymStack.size() > 0) {
                restoreState(this.tokenState);          // 将状态还原为上一个元素的状态
                cta.setEmpty();
                cta.append(this.synonymStack.pop());    // 获取并追加同义词
                pia.setPositionIncrement(0);            // 设置位置增量为0
                return true;
            }
            if (input.incrementToken()) {
                // 注意:当发现当前元素存在同义词之后,不能立即追加同义词,即不能在目标元素上直接处理
                if (this.isHaveSynonym(cta.toString())) {
                    this.tokenState = captureState();   // 存在同义词时,则捕获并保存当前状态
                }
                return true;
            } else {
                return false; // 只要TokenStream中没有元素,就返回false
            }
        }
    }

    /**
     * 单元测试要求MMSegTokenizer为final的
     */
    final class MyMMSegTokenizer extends MMSegTokenizer {

        public MyMMSegTokenizer(Seg seg, Reader reader) {
            super(seg, reader);
        }
    }
}

